#!/usr/bin/env python
# coding: utf-8

# In[4]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tushare as ts
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import warnings
warnings.filterwarnings('ignore')

class DataLoader:
    def __init__(self, start_date, end_date, token='your_tushare_token'):
        self.start_date = start_date
        self.end_date = end_date
        self.token = token
        # 初始化Tushare
        ts.set_token(self.token)
        self.pro = ts.pro_api()
    
    def load_stock_data(self, ticker):
        """使用Tushare获取股票数据"""
        try:
            # 股票日线行情
            df = ts.pro_bar(ts_code=ticker, adj='qfq', 
                           start_date=self.start_date.replace('-', ''),
                           end_date=self.end_date.replace('-', ''))
            
            if df is None or df.empty:
                print(f"无法获取{ticker}的数据")
                return None
            
            # 调整列名以匹配原代码
            df = df.rename(columns={
                'ts_code': 'Ticker',
                'trade_date': 'Date',
                'open': 'Open',
                'high': 'High',
                'low': 'Low',
                'close': 'Close',
                'vol': 'Volume'
            })
            
            # 转换日期格式
            df['Date'] = pd.to_datetime(df['Date'])
            df.set_index('Date', inplace=True)
            df.sort_index(inplace=True)
            
            return df
        except Exception as e:
            print(f"获取{ticker}数据时出错: {e}")
            return None
    
    def load_multiple_stocks(self, tickers):
        all_data = {}
        for ticker in tickers:
            data = self.load_stock_data(ticker)
            if data is not None:
                all_data[ticker] = data
        return all_data
    
    def get_fundamental_data(self, ticker):
        """获取基本面数据（PE、PB、ROE等）"""
        try:
            # 获取财务指标数据
            df = self.pro.fina_indicator(ts_code=ticker, 
                                        start_date=self.start_date.replace('-', ''),
                                        end_date=self.end_date.replace('-', ''))
            
            if df is None or df.empty:
                print(f"无法获取{ticker}的基本面数据")
                return None
            
            # 调整列名
            df = df.rename(columns={
                'ts_code': 'Ticker',
                'end_date': 'Date',
                'pe': 'PE',          # 市盈率
                'pb': 'PB',          # 市净率
                'roe': 'ROE',        # 净资产收益率
                'gross_margin': 'GM' # 毛利率
            })
            
            # 转换日期格式
            df['Date'] = pd.to_datetime(df['Date'])
            df.set_index('Date', inplace=True)
            df.sort_index(inplace=True)
            
            return df[['PE', 'PB', 'ROE', 'GM']]
        except Exception as e:
            print(f"获取{ticker}基本面数据时出错: {e}")
            return None

class TechnicalStrategy:
    def __init__(self, data):
        self.data = data.copy()
    
    def generate_signals(self):
        # 计算移动平均线
        self.data['SMA5'] = self.data['Close'].rolling(window=5).mean()
        self.data['SMA20'] = self.data['Close'].rolling(window=20).mean()
        
        # 计算RSI
        delta = self.data['Close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        self.data['RSI'] = 100 - (100 / (1 + rs))
        
        # 计算MACD
        self.data['EMA12'] = self.data['Close'].ewm(span=12, adjust=False).mean()
        self.data['EMA26'] = self.data['Close'].ewm(span=26, adjust=False).mean()
        self.data['MACD'] = self.data['EMA12'] - self.data['EMA26']
        self.data['Signal'] = self.data['MACD'].ewm(span=9, adjust=False).mean()
        self.data['Histogram'] = self.data['MACD'] - self.data['Signal']
        
        # 生成交易信号
        self.data['Signal'] = 0
        # 买入信号：短期均线上穿长期均线且RSI<30且MACD金叉
        self.data.loc[(self.data['SMA5'] > self.data['SMA20']) & 
                     (self.data['SMA5'].shift(1) <= self.data['SMA20'].shift(1)) &
                     (self.data['RSI'] < 30) &
                     (self.data['Histogram'] > 0) &
                     (self.data['Histogram'].shift(1) <= 0), 'Signal'] = 1
        
        # 卖出信号：短期均线下穿长期均线且RSI>70且MACD死叉
        self.data.loc[(self.data['SMA5'] < self.data['SMA20']) & 
                     (self.data['SMA5'].shift(1) >= self.data['SMA20'].shift(1)) &
                     (self.data['RSI'] > 70) &
                     (self.data['Histogram'] < 0) &
                     (self.data['Histogram'].shift(1) >= 0), 'Signal'] = -1
        
        return self.data

class FundamentalStrategy:
    def __init__(self, data, fundamental_data):
        self.data = data.copy()
        self.fundamental_data = fundamental_data
    
    def generate_signals(self):
        # 合并价格数据和基本面数据
        merged_data = pd.merge(self.data, self.fundamental_data, 
                              left_index=True, right_index=True, how='left')
        
        # 填充缺失值
        merged_data.fillna(method='ffill', inplace=True)
        
        # 基于基本面指标生成信号
        merged_data['Signal'] = 0
        
        # 买入信号：低PE、低PB、高ROE
        merged_data.loc[(merged_data['PE'] < merged_data['PE'].quantile(0.25)) &
                      (merged_data['PB'] < merged_data['PB'].quantile(0.25)) &
                      (merged_data['ROE'] > merged_data['ROE'].quantile(0.75)), 'Signal'] = 1
        
        # 卖出信号：高PE、高PB、低ROE
        merged_data.loc[(merged_data['PE'] > merged_data['PE'].quantile(0.75)) &
                      (merged_data['PB'] > merged_data['PB'].quantile(0.75)) &
                      (merged_data['ROE'] < merged_data['ROE'].quantile(0.25)), 'Signal'] = -1
        
        return merged_data

class MLStrategy:
    def __init__(self, data):
        self.data = data.copy()
    
    def prepare_features(self):
        # 计算技术指标
        self.data['SMA5'] = self.data['Close'].rolling(window=5).mean()
        self.data['SMA20'] = self.data['Close'].rolling(window=20).mean()
        self.data['SMA50'] = self.data['Close'].rolling(window=50).mean()
        
        self.data['EMA12'] = self.data['Close'].ewm(span=12, adjust=False).mean()
        self.data['EMA26'] = self.data['Close'].ewm(span=26, adjust=False).mean()
        
        self.data['RSI'] = self.calculate_rsi()
        
        # 计算价格变化率
        self.data['Return_1d'] = self.data['Close'].pct_change(1)
        self.data['Return_5d'] = self.data['Close'].pct_change(5)
        
        # 计算波动率
        self.data['Volatility_10d'] = self.data['Return_1d'].rolling(window=10).std()
        
        # 计算成交量变化
        self.data['Volume_Change'] = self.data['Volume'].pct_change()
        
        # 移除缺失值
        self.data.dropna(inplace=True)
        
        return self.data
    
    def calculate_rsi(self, period=14):
        delta = self.data['Close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
        return rsi
    
    def generate_signals(self):
        # 准备特征
        data = self.prepare_features()
        
        # 创建目标变量：未来5天收益率为正标记为1，否则为0
        data['Target'] = np.where(data['Close'].shift(-5) > data['Close'], 1, 0)
        
        # 移除缺失值
        data.dropna(inplace=True)
        
        # 分割特征和目标
        features = ['SMA5', 'SMA20', 'SMA50', 'EMA12', 'EMA26', 'RSI', 
                   'Return_1d', 'Return_5d', 'Volatility_10d', 'Volume_Change']
        X = data[features]
        y = data['Target']
        
        # 分割训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
        
        # 训练随机森林分类器
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model.fit(X_train, y_train)
        
        # 在测试集上预测
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"模型准确率: {accuracy:.4f}")
        
        # 将预测结果添加到数据中
        data['Prediction'] = 0
        data.loc[X_test.index, 'Prediction'] = y_pred
        
        # 生成交易信号
        data['Signal'] = 0
        data.loc[data['Prediction'] == 1, 'Signal'] = 1  # 买入信号
        data.loc[data['Prediction'] == 0, 'Signal'] = -1  # 卖出信号
        
        return data

class TradingSystem:
    def __init__(self, data, initial_capital=100000):
        self.data = data.copy()
        self.initial_capital = initial_capital
        self.positions = 0
        self.capital = initial_capital
        self.portfolio = pd.DataFrame(index=data.index)
        self.portfolio['positions'] = 0
        self.portfolio['capital'] = initial_capital
        self.portfolio['total'] = initial_capital
    
    def backtest(self, stop_loss=0.08, take_profit=0.15):
        signals = self.data['Signal']
        close_prices = self.data['Close']
        
        for i in range(1, len(self.data)):
            # 当前持仓和资金
            current_positions = self.portfolio['positions'].iloc[i-1]
            current_capital = self.portfolio['capital'].iloc[i-1]
            
            # 检查止损和止盈
            if current_positions > 0:
                buy_price = close_prices.iloc[i-1]
                current_price = close_prices.iloc[i]
                profit_pct = (current_price / buy_price) - 1
                
                # 止损触发
                if profit_pct < -stop_loss:
                    current_capital += current_positions * current_price
                    current_positions = 0
                
                # 止盈触发
                elif profit_pct > take_profit:
                    current_capital += current_positions * current_price
                    current_positions = 0
            
            # 执行交易信号
            if signals.iloc[i] == 1 and current_positions == 0:  # 买入信号
                # 用80%的资金买入
                shares_to_buy = int((current_capital * 0.8) / close_prices.iloc[i])
                current_positions += shares_to_buy
                current_capital -= shares_to_buy * close_prices.iloc[i]
            
            elif signals.iloc[i] == -1 and current_positions > 0:  # 卖出信号
                current_capital += current_positions * close_prices.iloc[i]
                current_positions = 0
            
            # 更新投资组合
            self.portfolio['positions'].iloc[i] = current_positions
            self.portfolio['capital'].iloc[i] = current_capital
            self.portfolio['total'].iloc[i] = current_capital + current_positions * close_prices.iloc[i]
        
        # 计算收益率
        self.calculate_performance()
        return self.portfolio
    
    def calculate_performance(self):
        # 计算每日收益率
        self.portfolio['daily_return'] = self.portfolio['total'].pct_change()
        
        # 计算累积收益率
        self.portfolio['cumulative_return'] = (1 + self.portfolio['daily_return']).cumprod() - 1
        
        # 计算最大回撤
        self.portfolio['peak'] = self.portfolio['total'].cummax()
        self.portfolio['drawdown'] = (self.portfolio['total'] / self.portfolio['peak']) - 1
        max_drawdown = self.portfolio['drawdown'].min()
        
        # 计算年化收益率
        days = len(self.portfolio)
        annual_return = ((1 + self.portfolio['cumulative_return'].iloc[-1]) ** (252 / days)) - 1
        
        # 计算夏普比率
        risk_free_rate = 0.03  # 假设无风险利率为3%
        daily_risk_free = (1 + risk_free_rate) ** (1/252) - 1
        excess_return = self.portfolio['daily_return'] - daily_risk_free
        sharpe_ratio = np.sqrt(252) * excess_return.mean() / excess_return.std()
        
        print(f"总收益率: {self.portfolio['cumulative_return'].iloc[-1]:.2%}")
        print(f"年化收益率: {annual_return:.2%}")
        print(f"最大回撤: {max_drawdown:.2%}")
        print(f"夏普比率: {sharpe_ratio:.2f}")
        
        return {
            'total_return': self.portfolio['cumulative_return'].iloc[-1],
            'annual_return': annual_return,
            'max_drawdown': max_drawdown,
            'sharpe_ratio': sharpe_ratio
        }
    
    def plot_performance(self):
        plt.figure(figsize=(12, 8))
        
        # 绘制资产曲线
        plt.subplot(2, 1, 1)
        plt.plot(self.portfolio['total'])
        plt.title('Portfolio Value')
        plt.grid(True)
        
        # 绘制回撤曲线
        plt.subplot(2, 1, 2)
        plt.fill_between(self.portfolio.index, self.portfolio['drawdown'], 0, color='red', alpha=0.3)
        plt.title('Drawdown')
        plt.grid(True)
        
        plt.tight_layout()
        plt.show()

def run_backtest(ticker, start_date, end_date, token):
    print(f"\n回测 {ticker} 从 {start_date} 到 {end_date}")
    
    # 加载数据
    data_loader = DataLoader(start_date, end_date, token)
    data = data_loader.load_stock_data(ticker)
    
    if data is None:
        return None
    
    # 1. 技术策略回测
    print("\n=== 技术策略回测 ===")
    tech_strategy = TechnicalStrategy(data)
    tech_data = tech_strategy.generate_signals()
    tech_system = TradingSystem(tech_data)
    tech_portfolio = tech_system.backtest()
    tech_performance = tech_system.calculate_performance()
    
    # 2. 获取基本面数据并回测
    print("\n=== 基本面策略回测 ===")
    fundamental_data = data_loader.get_fundamental_data(ticker)
    if fundamental_data is not None:
        fund_strategy = FundamentalStrategy(data, fundamental_data)
        fund_data = fund_strategy.generate_signals()
        fund_system = TradingSystem(fund_data)
        fund_portfolio = fund_system.backtest()
        fund_performance = fund_system.calculate_performance()
    else:
        print("无法获取基本面数据，跳过基本面策略回测")
        fund_performance = None
    
    # 3. 机器学习策略回测
    print("\n=== 机器学习策略回测 ===")
    ml_strategy = MLStrategy(data)
    ml_data = ml_strategy.generate_signals()
    ml_system = TradingSystem(ml_data)
    ml_portfolio = ml_system.backtest()
    ml_performance = ml_system.calculate_performance()
    
    # 比较策略
    print("\n=== 策略比较 ===")
    print(f"技术策略总收益率: {tech_performance['total_return']:.2%}")
    if fund_performance:
        print(f"基本面策略总收益率: {fund_performance['total_return']:.2%}")
    print(f"机器学习策略总收益率: {ml_performance['total_return']:.2%}")
    
    # 绘制绩效对比图
    plt.figure(figsize=(12, 6))
    plt.plot(tech_portfolio['cumulative_return'], label='技术策略')
    if fund_performance:
        plt.plot(fund_portfolio['cumulative_return'], label='基本面策略')
    plt.plot(ml_portfolio['cumulative_return'], label='机器学习策略')
    plt.title(f'{ticker} 不同策略累积收益率对比')
    plt.xlabel('日期')
    plt.ylabel('累积收益率')
    plt.legend()
    plt.grid(True)
    plt.show()
    
    return {
        'technical': tech_performance,
        'fundamental': fund_performance,
        'machine_learning': ml_performance
    }

if __name__ == "__main__":
    # 请替换为你的Tushare token
    TUSHARE_TOKEN = '7705d187ddde476492fc7c72b674218f8d383448f7126a0457392231'
    
    # 回测参数
    start_date = "2020-01-01"
    end_date = "2023-12-31"
    # 使用Tushare格式的股票代码
    tickers = ["000001.SZ", "600519.SH", "002594.SZ", "601318.SH", "600036.SH"]
    
    # 结果汇总
    results = {}
    
    # 回测多支股票
    for ticker in tickers:
        result = run_backtest(ticker, start_date, end_date, TUSHARE_TOKEN)
        if result:
            results[ticker] = result
    
    # 显示总体结果
    print("\n=== 多股票回测汇总 ===")
    for ticker, result in results.items():
        print(f"\n{ticker}:")
        print(f"  技术策略年化收益率: {result['technical']['annual_return']:.2%}")
        if result['fundamental']:
            print(f"  基本面策略年化收益率: {result['fundamental']['annual_return']:.2%}")
        print(f"  机器学习策略年化收益率: {result['machine_learning']['annual_return']:.2%}")      


# In[ ]:




